package com.raos.example.rag;

import com.raos.example.service.Assistant;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.List;


/**
 * RAG聊天 webAPI
 *
 * @author raos
 * @date 2025-03-20
 */
@Slf4j
@RequestMapping("/api/rag")
@RestController
public class RagAPI {

    @Resource(name = "assistant")
    Assistant assistant;

    @Resource(name = "inMemoryEmbeddingStore")
    EmbeddingStore<TextSegment> embeddingStore;

    @GetMapping("/high/chat")
    public String highChat(@RequestParam(value = "message") String message) {
        log.info("message=[{}]", message);
        String chat = assistant.chat(message);
        log.info("chat=[{}]", chat);
        return chat;
    }

/*
提问：请问你知道行明吗？
AI回答（通过额外检索能力）：
⾏明是明朝末年⼀位颇具传奇⾊彩的侠客，他不仅以⾏侠仗义闻名，更以其幽默与机智的⻛采赢得了江湖⼈⼠的尊敬和喜爱。他的故事常常展现出他如何⽤轻松诙谐的⽅式化解纷争，⽽⾮单纯依靠武⼒。

例如，在⼀次事件中，⾏明遇到了⼀群武⼠围堵⼀位年轻姑娘，误以为她偷了宝物。⾏明并未直接动武，⽽是以幽默的⾔辞和机智的提议，带着姑娘和武⼠们⼀同回到府邸查明真相。途中，他以玩笑缓和了紧张的⽓氛，最终发现所谓的“宝物”不过是⼀个普通的⽟佩。⾏明以轻松的态度化解了误会，并让武⼠们意识到⾃⼰的错误。

⾏明的⾔⾏让⼈们看到，江湖并⾮只有⼑光剑影，侠义之道也可以通过幽默与智慧来实现。正如他所说：“江湖不⼀定总是⼑光剑影，偶尔⼀场笑谈，亦能化解许多误会。” 他的故事也因此成为了江湖中⼝⼝相传的佳话。
*/

    /**
     * 加载文档至内存向量库（开启内存向量库）
     */
    @PostMapping("/loadInMemory")
    public String load() {
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("D:\\Users\\HASEE\\ai-demo\\lang-chain4j-demo\\documents");
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);
        log.info("loadInMemory success");
        return "success";
    }

    @Resource(name = "assistant2")
    Assistant assistant2;

    @Resource(name = "pgVectorEmbeddingStore")
    EmbeddingStore<TextSegment> embeddingStore2;

    @Autowired
    EmbeddingModel embeddingModel;

    @GetMapping("/high/chat2")
    public String highChat2(@RequestParam(value = "message") String message) {
        log.info("message=[{}]", message);
        String chat = assistant2.chat(message);
        log.info("chat=[{}]", chat);
        return chat;
    }

    /**
     * 加载文档至PgVector向量库
     */
    @PostMapping("/load2PgVector")
    public String load2() {
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("D:\\Users\\HASEE\\ai-demo\\lang-chain4j-demo\\documents");
        // 默认文本切分
        // EmbeddingStoreIngestor.ingest(documents, embeddingStore2);
        // 自定义文本切分（注意文本切分要合理）
        EmbeddingStoreIngestor.builder().embeddingStore(embeddingStore2)
                .embeddingModel(embeddingModel)
                .documentSplitter(new DocumentByLineSplitter(200, 50))
                .build().ingest(documents);
        log.info("loadInMemory success");
        return "success";
    }

}
